import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import PchipInterpolator


def draw_pc(result_dict, option, data, label):
    global np
    """
    路径可视化与后处理函数
    参数:
        result: 包含路径信息的结果字典
        option: 算法参数配置
        data: 环境数据字典
        str_label: 算法标签字符串
    返回:
        optimized_path: 优化后的路径点数组
    """
    path0 = []
    # 强制添加起点到路径数组
    start_point = data['S'][:3]
    path0.append(start_point)

    # 提取并处理原始路径
    result_path_ndarray = result_dict["path"][0]
    now_path_index = 0
    aim_path_index = 1

    while aim_path_index < result_path_ndarray.shape[0]:
        pass

        # 检查路径段是否与地形碰撞
        # 修正列索引以匹配三维坐标数据结构
        # 验证数组结构并提取坐标
        if result_path_ndarray.shape[1] >= 3:
            now_path = result_path_ndarray[now_path_index, :3].flatten()  # 获取前3列
            aim_path = result_path_ndarray[aim_path_index, :3].flatten()  # 获取前3列
        else:
            # 处理数组结构异常
            now_path = np.array([0, 0, 0])
            aim_path = np.array([0, 0, 0])
        collision_flag = 0

        # 采样检查路径段
        for i in range(1000):
            # 显式计算三维坐标插值
            x = now_path[0] + (aim_path[0] - now_path[0]) * i / 1000
            y = now_path[1] + (aim_path[1] - now_path[1]) * i / 1000
            h = now_path[2] + (aim_path[2] - now_path[2]) * i / 1000
            interp_point = np.array([x, y, h])
            x0, y0 = round(interp_point[0]), round(interp_point[1])

            # 边界检查
            x0 = np.clip(x0, 1, data['map_z'].shape[0])
            y0 = np.clip(y0, 1, data['map_z'].shape[1])

            # 获取地形高度
            # 修复坐标顺序以正确获取地形高度
            terrain_h = data['map_z'][y0-1, x0-1]  # MATLAB索引转Python
            if terrain_h > interp_point[2]:
                collision_flag = 1
                break

        # 递增索引以处理下一段路径
        now_path_index += 1
        aim_path_index += 1

        if collision_flag == 1:
            path0.append(now_path)
            path0.append(result_path_ndarray[aim_path_index - 2, 0:3])
            now_path_index = aim_path_index
            aim_path_index = now_path_index + 1
        else:
                now_path_index += 1
                aim_path_index += 1

        path0.append(now_path)
        path0.append(aim_path)
    # end while

    # 转换为numpy数组
    path0 = np.vstack(path0) if path0 else np.empty((0, 3))
    # 去重连续重复点
    if len(path0) > 1:
        mask = np.any(np.diff(path0, axis=0) != 0, axis=1)
        mask = np.concatenate([[True], mask])
        path0 = path0[mask]
    optimized_path = []

    # 路径平滑处理
    if path0.shape[0] >= 2:
        tempN = path0[:2, :].copy()
        nextInd = 2

        # 方向向量分析
        flag = np.ones(2)
        type_dir = np.zeros(2)
        for j in range(2):
            if np.isclose(tempN[0, j], tempN[1, j]):
                flag[j] = 0
            type_dir[j] = 1 if tempN[0, j] < tempN[1, j] else 2

        # 路径点插值
        while nextInd < path0.shape[0]:
            flag0 = flag.copy()
            for j in range(2):
                if type_dir[j] == 1 and tempN[-1, j] >= path0[nextInd, j]:
                    flag0[j] = 0
                elif type_dir[j] == 2 and tempN[-1, j] <= path0[nextInd, j]:
                    flag0[j] = 0

            if np.sum(flag0) == 0:
                # 执行插值
                if tempN.shape[0] > 2:
                    tempN[-1] = tempN[-2] + 0.5 * (tempN[-1] - tempN[-2])

                # 选择插值维度
                x, y, z = tempN[:, 0], tempN[:, 1], tempN[:, 2]

                # 插值前处理x值
                order = np.argsort(x)
                x, y, z = x[order], y[order], z[order]

                # 去除重复的x值
                _, idx = np.unique(x, return_index=True)
                x, y, z = x[idx], y[idx], z[idx]

                # 确保x值严格递增
                if len(x) >= 2 and np.all(np.diff(x) > 0):
                    xi = np.linspace(x[0],x[1], 100)
                    yi = PchipInterpolator(x, y)(xi)
                    zi = PchipInterpolator(x, z)(xi)
                else:
                    xi = np.linspace(tempN[0, 0], tempN[-1, 0], 100)
                    yi = np.interp(xi, x, y)
                    zi = np.interp(xi, x, z)

                optimized_path.extend(np.column_stack((xi, yi, zi)))

                # 更新路径点
                path0 = path0[nextInd-1:]
                tempN = path0[:2, :].copy()
                nextInd = 2
                flag = np.ones(2)
                for j in range(2):
                    if np.isclose(tempN[0, j], tempN[1, j]):
                        flag[j] = 0
                    type_dir[j] = 1 if tempN[0, j] < tempN[1, j] else 2
            else:
                tempN = np.vstack((tempN, path0[nextInd]))
                nextInd += 1

        # 处理剩余点
        # 选择插值维度
        x, y, z = tempN[:, 0], tempN[:, 1], tempN[:, 2]

        # 插值前处理x值
        order = np.argsort(x)
        x, y, z = x[order], y[order], z[order]

        # 去除重复的x值
        _, idx = np.unique(x, return_index=True)
        x, y, z = x[idx], y[idx], z[idx]

        # 确保x值严格递增
        if len(x) >= 2 and np.all(np.diff(x) > 0):
            xi = np.linspace(x[0], x[1], 100)
            yi = PchipInterpolator(x, y)(xi)
            zi = PchipInterpolator(x, z)(xi)
        else:
            xi = np.linspace(tempN[0, 0], tempN[-1, 0], 100)
            yi = np.interp(xi, x, y)
            zi = np.interp(xi, x, z)

        optimized_path.extend(np.column_stack((xi, yi, zi)))

    # 转换为numpy数组并确保安全高度
    optimized_path = np.vstack(optimized_path).reshape(-1, 3) if optimized_path else np.empty((0, 3))
    if optimized_path.size > 0:
        # 向量化计算地形安全高度
        x0 = np.clip(np.round(optimized_path[:, 0]).astype(int), 1, data['map_z'].shape[0])
        y0 = np.clip(np.round(optimized_path[:, 1]).astype(int), 1, data['map_z'].shape[1])
        terrain_h = data['map_z'][x0-1, y0-1] + 1
        optimized_path[:, 2] = np.maximum(optimized_path[:, 2], terrain_h)

    # 路径优化: 简化 + 平滑
    from scipy.interpolate import make_interp_spline
    from scipy.spatial.distance import cdist
    
    # Ramer-Douglas-Peucker路径简化算法
    def simplify_path(points, epsilon=1.0):
        if len(points) <= 2:
            return points
        
        # 计算最远点
        start, end = points[0], points[-1]
        distances = cdist(points[1:-1], [start, end]).min(axis=1)
        max_dist_idx = np.argmax(distances) + 1
        
        # 递归简化
        if distances[max_dist_idx - 1] > epsilon:
            left = simplify_path(points[:max_dist_idx+1], epsilon)
            right = simplify_path(points[max_dist_idx:], epsilon)
            return np.vstack((left[:-1], right))
        return np.vstack((start, end))
    
    # 简化路径
    # 增加epsilon值增强路径简化效果
    # 进一步增加epsilon值以增强路径简化
    simplified = simplify_path(optimized_path, epsilon=10.0)
    # 确保路径以起点开始
    start_point = data['S'][:3]
    if len(simplified) == 0 or not np.allclose(simplified[0], start_point):
        simplified = np.vstack((start_point, simplified))
    
    # 确保路径至少有3个点用于平滑
    if len(simplified) < 3:
        simplified = optimized_path
        # 回退时仍需确保起点存在
        if len(simplified) == 0 or not np.allclose(simplified[0], start_point):
            simplified = np.vstack((start_point, simplified))
    
    # 提取坐标
    x, y, z = simplified[:, 0], simplified[:, 1], simplified[:, 2]
    
    # 三次样条平滑
    t = np.arange(len(x))
    t_new = np.linspace(t.min(), t.max(), 200)
    
    spl_x = make_interp_spline(t, x, k=3)
    spl_y = make_interp_spline(t, y, k=3)
    spl_z = make_interp_spline(t, z, k=3)
    
    # 生成平滑路径
    smoothed_path = np.column_stack((
        spl_x(t_new),
        spl_y(t_new),
        spl_z(t_new)
    ))
    
    # 应用安全高度调整到平滑路径
    if smoothed_path.size > 0:
        x0 = np.clip(np.round(smoothed_path[:, 0]).astype(int), 1, data['map_z'].shape[0])
        y0 = np.clip(np.round(smoothed_path[:, 1]).astype(int), 1, data['map_z'].shape[1])
        terrain_h = data['map_z'][x0-1, y0-1] + 1
        smoothed_path[:, 2] = np.maximum(smoothed_path[:, 2], terrain_h)
    
    # 绘制地形
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    surface = ax.plot_surface(data['map_x'], data['map_y'], data['map_z'], cmap='terrain', alpha=0.7)
    fig.colorbar(surface, ax=ax, shrink=0.5, aspect=5)

    # 标记起点和终点
    print("标记起点和终点")
    ax.scatter(data['S'][1], data['S'][0], data['S'][2], c='g', s=100, marker='o', label='起点')
    ax.scatter(data['E0'][:, 1], data['E0'][:, 0], data['E0'][:, 2], c='r', s=100, marker='^', label='终点')

    # 绘制路径
    print("正在绘制路径...")
    if smoothed_path.ndim == 1:
        smoothed_path = smoothed_path.reshape(-1, 3)
    print("优化后路径点信息：", smoothed_path)
    ax.plot(simplified[:, 0], simplified[:, 1], simplified[:, 2], 'r-', linewidth=5)

    # 设置标签和标题
    print("设置成图标签和标题")
    ax.set_title(f'{label}目标值: {result_dict["fit"]:.4f}')
    ax.set_xlabel('X坐标')
    ax.set_ylabel('Y坐标')
    ax.set_zlabel('高度')
    ax.legend()
    ax.grid(True)

    print("正在渲染路径图...")
    plt.tight_layout()
    plt.show()
    
    return simplified  # 返回包含起点的简化路径